
#include "BasicClient.h"
#include "opencv2/opencv.hpp"



using  namespace std;

const uint32_t g_modelOutputBoxNum_v5 = 25200;    //20160   80640
const uint32_t g_modelOutputBoxNum_v8 = 8400;

const uint32_t g_all_v5 = 85;    //5 + 80
const uint32_t g_all_v8 = 84;    //4 + 80
const double g_NMSThreshold = 0.45;
const double g_scoreThreshold = 0.25;
const uint32_t g_modelWidth = 640;
const uint32_t g_modelHeight = 640;


bool SortScore(BBox box1, BBox box2)
{
    return box1.score > box2.score;
}

float IOU(const BBox &b1, const BBox &b2)
{
    float x1 = max(b1.x - b1.w / 2, b2.x - b2.w / 2);
    float y1 = max(b1.y - b1.h / 2, b2.y - b2.h / 2);
    float x2 = min(b1.x + b1.w / 2, b2.x + b2.w / 2);
    float y2 = min(b1.y + b1.h / 2, b2.y + b2.h / 2);
    float w = max(0.0f, x2 - x1 + 1);
    float h = max(0.0f, y2 - y1 + 1);
    float area = w * h;
    return area / (b1.w * b1.h + b2.w * b2.h - area);
}


void NMS(vector<BBox> &boxes, vector<BBox> &result)
{
    result.clear();
    sort(boxes.begin(), boxes.end(), SortScore);  //按照置信度从大到小排序

    while (boxes.size() != 0) {
        result.push_back(boxes[0]);
        size_t index = 1;
        while (boxes.size() > index) {
            float iou = IOU(boxes[0], boxes[index]);
            if (iou > g_NMSThreshold) {
                boxes.erase(boxes.begin() + index);
                continue;
            }
            ++index;
            //cout<<index<<endl;
        }
        boxes.erase(boxes.begin());
    }
}


void DrawBoundBoxToImage(vector<BBox>& detectionResults, cv::Mat& image)
{
    const vector<cv::Scalar> g_colors {
            cv::Scalar(0, 0, 255), cv::Scalar(255, 0, 0), cv::Scalar(0, 125, 255)
    };

    int half = 2;
    float w = float(image.cols) / float(g_modelWidth);
    float h = float(image.rows) / float(g_modelHeight);
    for (int i = 0; i < detectionResults.size(); ++i) {
        cv::Point p1, p2;
        p1.x = (uint32_t)((detectionResults[i].x - detectionResults[i].w / half) * w);
        p1.y = (uint32_t)((detectionResults[i].y - detectionResults[i].h / half) * h);
        p2.x = (uint32_t)((detectionResults[i].x + detectionResults[i].w / half) * w);
        p2.y = (uint32_t)((detectionResults[i].y + detectionResults[i].h / half) * h);
        cv::rectangle(image, p1, p2, g_colors[detectionResults[i].classIndex], 2);
        // char name[50];
        // sprintf(name, "%.1d_%.3f", detectionResults[i].classIndex, detectionResults[i].score);
        // // 图片，添加的文字，左上角坐标，字体，字体大小，颜色，字体粗细
        // cv::putText(image, name, p1, cv::FONT_HERSHEY_SIMPLEX, 1, g_colors[detectionResults[i].classIndex], 1);
    }

    // cv::imshow("output_pic", image);
    cv::imwrite("./output_pic.jpg", image);
}


int main() {

    BasicClient basicClient("yolov8n.om",g_modelWidth,g_modelHeight);
    basicClient.Init_atlas();

    basicClient.CreateInput();
    basicClient.CreateOutput();
    cv::Mat img = cv::imread("./1.jpg");
    cv::Mat frame;
    cv::resize(img, frame, cv::Size(g_modelWidth, g_modelHeight));
    aclError ret = aclrtMemcpy(basicClient.g_imageDataBuf_, basicClient.g_imageDataSize_,
                               frame.ptr<uint8_t>(), basicClient.g_imageDataSize_, ACL_MEMCPY_HOST_TO_DEVICE);
    if (ret) {
        std::cout<<"Copy resized image data to device failed."<<std::endl;
        return FAILED;
    }
    basicClient.inference();
    uint32_t dataSize = 0;
    float* detectData = (float*)basicClient.GetInferenceOutputItem(dataSize, 0);
    vector<BBox> boxes;
    for (uint32_t i = 0; i < g_modelOutputBoxNum_v8; i++) {
        float  maxScore = 0.0f;
        int    maxClass = -1;
        for (uint32_t j = 4; j < g_all_v8; ++j) {
            float score = detectData[j * g_modelOutputBoxNum_v8 + i];
            if (score > maxScore) {
                maxScore = score;
                maxClass = int(j - 4);
            }
        }

        if (maxScore >= g_scoreThreshold) {
            // 读取框的四个坐标：x, y, w, h
            BBox b;
            b.x          = detectData[0 * g_modelOutputBoxNum_v8 + i];
            b.y          = detectData[1 * g_modelOutputBoxNum_v8 + i];
            b.w          = detectData[2 * g_modelOutputBoxNum_v8 + i];
            b.h          = detectData[3 * g_modelOutputBoxNum_v8 + i];
            b.score      = maxScore;
            b.classIndex = maxClass;
            b.index      = i;
            boxes.push_back(b);
        }
    }

    // 执行 NMS
    vector<BBox> result;
    NMS(boxes, result);

    std::cout<<"检测个数："<<result.size()<<std::endl;
    DrawBoundBoxToImage(result, img);
    if (basicClient.g_runMode_ == ACL_HOST) {
        aclError ret = aclrtFreeHost(detectData);
    }




    basicClient.DestroyResource();
    return 0;
}
